import math
import torch
import torch.nn as nn
import torch.nn.functional as F

def get_mask(input_size, window_set, inner_size, device):
    """Get the attention mask of PAM-Naive"""
    # Get the size of all layers
    all_size = []
    all_size.append(input_size)
    for i in range(len(window_set)):
        layer_size = math.floor(all_size[i] / window_set[i])
        all_size.append(layer_size)

    seq_length = sum(all_size)
    mask = torch.zeros(seq_length, seq_length, device=device)

    # get intra-scale mask
    inner_window = inner_size // 2
    for layer_idx in range(len(all_size)):
        start = sum(all_size[:layer_idx])
        for i in range(start, start + all_size[layer_idx]):
            left_side = max(i - inner_window, start)
            right_side = min(i + inner_window + 1, start + all_size[layer_idx])
            mask[i, left_side:right_side] = 1

    # get inter-scale mask
    for layer_idx in range(1, len(all_size)):
        start = sum(all_size[:layer_idx])
        for i in range(start, start + all_size[layer_idx]):
            left_side = (start - all_size[layer_idx - 1]) + (i - start) * window_set[layer_idx - 1]
            if i == ( start + all_size[layer_idx] - 1):
                right_side = start
            else:
                right_side = (start - all_size[layer_idx - 1]) + (i - start + 1) * window_set[layer_idx - 1]
            mask[i, left_side:right_side] = 1
            mask[left_side:right_side, i] = 1

    mask = (1 - mask).bool()

    return mask, all_size


def refer_points(all_sizes, window_set, device):
    """Gather features from PAM's pyramid sequences"""
    input_size = all_sizes[0]
    indexes = torch.zeros(input_size, len(all_sizes), device=device)

    for i in range(input_size):
        indexes[i][0] = i
        former_index = i
        for j in range(1, len(all_sizes)):
            start = sum(all_sizes[:j])
            inner_layer_idx = former_index - (start - all_sizes[j - 1])
            former_index = start + min(inner_layer_idx // window_set[j - 1], all_sizes[j] - 1)
            indexes[i][j] = former_index

    indexes = indexes.unsqueeze(0).unsqueeze(3)

    return indexes.long()

def last_points(all_sizes, device):
    layer_num = len(all_sizes)
    indexes = torch.zeros(1, layer_num, 1, device=device)
    for i in range(layer_num):
        indexes[0][i][0] = all_sizes[i] - 1
    return indexes.long()


class ScaledDotProductAttention(nn.Module):
    """ Scaled Dot-Product Attention """

    def __init__(self, temperature, attn_dropout=0.2):
        super().__init__()

        self.temperature = temperature
        self.dropout = nn.Dropout(attn_dropout)

    def forward(self, q, k, v, mask=None):
        attn = torch.matmul(q / self.temperature, k.transpose(2, 3))

        if mask is not None:
            attn = attn.masked_fill(mask, -1e9)

        attn = self.dropout(F.softmax(attn, dim=-1))
        output = torch.matmul(attn, v)

        return output, attn

class PositionalEmbedding(nn.Module):
    def __init__(self, d_model, max_len=5000):
        super(PositionalEmbedding, self).__init__()
        # Compute the positional encodings once in log space.
        pe = torch.zeros(max_len, d_model).float()
        pe.require_grad = False

        position = torch.arange(0, max_len).float().unsqueeze(1)
        div_term = (torch.arange(0, d_model, 2).float() * -(math.log(10000.0) / d_model)).exp()

        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)

        pe = pe.unsqueeze(0)
        self.register_buffer('pe', pe)

    def forward(self, x):
        return self.pe[:, :x.size(1)]


class MultiHeadAttention(nn.Module):
    """ Multi-Head Attention module """

    def __init__(self, n_head, d_model, d_k, d_v, dropout=0.1, normalize_before=True):
        super().__init__()

        self.normalize_before = normalize_before
        self.n_head = n_head
        self.d_k = d_k
        self.d_v = d_v

        self.w_qs = nn.Linear(d_model, n_head * d_k, bias=False)
        self.w_ks = nn.Linear(d_model, n_head * d_k, bias=False)
        self.w_vs = nn.Linear(d_model, n_head * d_v, bias=False)
        nn.init.xavier_uniform_(self.w_qs.weight)
        nn.init.xavier_uniform_(self.w_ks.weight)
        nn.init.xavier_uniform_(self.w_vs.weight)

        self.fc = nn.Linear(d_v * n_head, d_model)
        nn.init.xavier_uniform_(self.fc.weight)

        self.attention = ScaledDotProductAttention(temperature=d_k ** 0.5, attn_dropout=dropout)

        self.layer_norm = nn.LayerNorm(d_model, eps=1e-6)
        self.dropout = nn.Dropout(dropout)

    def forward(self, q, k, v, mask=None):
        d_k, d_v, n_head = self.d_k, self.d_v, self.n_head
        sz_b, len_q, len_k, len_v = q.size(0), q.size(1), k.size(1), v.size(1)

        residual = q
        if self.normalize_before:
            q = self.layer_norm(q)

        # Pass through the pre-attention projection: b x lq x (n*dv)
        # Separate different heads: b x lq x n x dv
        q = self.w_qs(q).view(sz_b, len_q, n_head, d_k)
        k = self.w_ks(k).view(sz_b, len_k, n_head, d_k)
        v = self.w_vs(v).view(sz_b, len_v, n_head, d_v)

        # Transpose for attention dot product: b x n x lq x dv
        q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)

        if mask is not None:
            if len(mask.size()) == 3:
                mask = mask.unsqueeze(1)  # For head axis broadcasting.

        output, attn = self.attention(q, k, v, mask=mask)

        # Transpose to move the head dimension back: b x lq x n x dv
        # Combine the last two dimensions to concatenate all the heads together: b x lq x (n*dv)
        output = output.transpose(1, 2).contiguous().view(sz_b, len_q, -1)
        output = self.dropout(self.fc(output))
        output += residual

        if not self.normalize_before:
            output = self.layer_norm(output)
        return output, attn

class PositionwiseFeedForward(nn.Module):
    """ Two-layer position-wise feed-forward neural network. """

    def __init__(self, d_in, d_hid, dropout=0.1, normalize_before=True):
        super().__init__()

        self.normalize_before = normalize_before

        self.w_1 = nn.Linear(d_in, d_hid)
        self.w_2 = nn.Linear(d_hid, d_in)

        self.layer_norm = nn.LayerNorm(d_in, eps=1e-6)
        #self.layer_norm = GraphNorm(d_in)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        residual = x
        if self.normalize_before:
            x = self.layer_norm(x)

        x = F.gelu(self.w_1(x))
        x = self.dropout(x)
        x = self.w_2(x)
        x = self.dropout(x)
        x = x + residual

        if not self.normalize_before:
            x = self.layer_norm(x)
        return x

class ConvLayer(nn.Module):
    def __init__(self, c_in, window_size):
        super(ConvLayer, self).__init__()
        self.downConv = nn.Conv1d(in_channels=c_in,
                                  out_channels=c_in,
                                  kernel_size=window_size,
                                  stride=window_size)
        self.norm = nn.BatchNorm1d(c_in)
        self.activation = nn.ELU()

    def forward(self, x):
        x = self.downConv(x)
        x = self.norm(x)
        x = self.activation(x)
        return x

    
class Bottleneck_Construct(nn.Module):
    """Bottleneck convolution CSCM"""
    def __init__(self, d_model, window_size, d_inner):
        super(Bottleneck_Construct, self).__init__()
        if not isinstance(window_size, list):
            self.conv_layers = nn.ModuleList([
                ConvLayer(d_inner, window_size),
                ConvLayer(d_inner, window_size),
                ConvLayer(d_inner, window_size)
                ])
        else:
            self.conv_layers = []
            for i in range(len(window_size)):
                self.conv_layers.append(ConvLayer(d_inner, window_size[i]))
            self.conv_layers = nn.ModuleList(self.conv_layers)
        self.up = nn.Linear(d_inner, d_model)
        self.down = nn.Linear(d_model, d_inner)
        self.norm = nn.LayerNorm(d_model)

    def forward(self, enc_input):

        temp_input = self.down(enc_input).permute(0, 2, 1)
        all_inputs = []
        for i in range(len(self.conv_layers)):
            temp_input = self.conv_layers[i](temp_input)
            all_inputs.append(temp_input)

        all_inputs = torch.cat(all_inputs, dim=2).transpose(1, 2)
        all_inputs = self.up(all_inputs)
        all_inputs = torch.cat([enc_input, all_inputs], dim=1)

        all_inputs = self.norm(all_inputs)

        return all_inputs

class SingleStepEmbedding(nn.Module):
    
    def __init__(self, feature_size, d_model, num_tokens, seq_length, device):
        super().__init__()
        self.k = d_model
        if num_tokens:
            self.token_in_count = True
            self.token_embedding = nn.Embedding(num_tokens, d_model)
        self.position = torch.arange(seq_length, device=device).unsqueeze(0)
        self.position_vec = torch.tensor([math.pow(10000.0, 2.0 * (i // 2) / d_model) for i in range(d_model)], device=device)

        self.linear = nn.Linear(feature_size, d_model)
        nn.init.xavier_normal_(self.linear.weight)
    
    def transformer_embedding(self, position, vector):
        """
        Input: batch*seq_len.
        Output: batch*seq_len*d_model.
        """
        result = position.unsqueeze(-1) / vector
        result[:, :, 0::2] = torch.sin(result[:, :, 0::2])
        result[:, :, 1::2] = torch.cos(result[:, :, 1::2])
        return result
    
    def forward(self, data, tokens):
        b, t, h = data.size()
        k = self.k
        
        sec = self.linear(data)
        if tokens is not None:
            token_emb = self.token_embedding(tokens)
            token_emb = token_emb.expand(t, b, k).transpose(0, 1)
            sec = sec + token_emb
        
        position = self.position.repeat(b, 1).to(data.device)
        position_emb = self.transformer_embedding(position, self.position_vec.to(data.device))
        sec = sec + position_emb
        return sec
        
        
        

class EncoderLayer(nn.Module):
    """ Compose with two layers """

    def __init__(self, d_model, d_inner, n_head, d_k, d_v, dropout=0.1, normalize_before=True):
        super().__init__()

        self.slf_attn = MultiHeadAttention(n_head, d_model, d_k, d_v, dropout=dropout, normalize_before=normalize_before)

        self.pos_ffn = PositionwiseFeedForward(
            d_model, d_inner, dropout=dropout, normalize_before=normalize_before)

    def forward(self, enc_input, slf_attn_mask=None):
        enc_output, enc_slf_attn = self.slf_attn(enc_input, enc_input, enc_input, mask=slf_attn_mask)

        enc_output = self.pos_ffn(enc_output)

        return enc_output, enc_slf_attn

class DecoderLayer(nn.Module):
    
    def __init__(self, d_model, d_inner, n_head, d_k, d_v, dropout=0.1, normalize_before=True):
        super().__init__()
        self.slf_attn = MultiHeadAttention(n_head, d_model, d_k, d_v, dropout=dropout, normalize_before=normalize_before)

        self.pos_ffn = PositionwiseFeedForward(
            d_model, d_inner, dropout=dropout, normalize_before=normalize_before)
    
    def forward(self, q, k, v, slf_attn_mask=None):
        enc_output, enc_slf_attn = self.slf_attn(q, k, v, mask=slf_attn_mask)
        enc_output = self.pos_ffn(enc_output)
        return enc_output, enc_slf_attn
        


class Encoder(nn.Module):
    
    def __init__(self, d_model, window_set, input_size, feature_size, n_layer, n_head, d_inner_hid, d_k, d_v, num_tokens, mask, dropout, device):
        super().__init__()
        
        self.d_model = d_model
        self.num_heads = n_head
        self.mask = mask
        
        self.layers = nn.ModuleList([
            EncoderLayer(d_model, d_inner_hid, n_head, d_k, d_v, dropout, normalize_before=False) for _ in range(n_layer)
        ])
        self.embedding = SingleStepEmbedding(feature_size, d_model, num_tokens, input_size, device)
        self.conv_layers = Bottleneck_Construct(d_model, window_set, d_k)
    
    def forward(self, sequence, token):
        seq_enc = self.embedding(sequence, token)
        mask = self.mask.repeat(len(seq_enc), self.num_heads, 1, 1).to(sequence.device)
        seq_enc = self.conv_layers(seq_enc)
        enc_seq = seq_enc
        for i in range(len(self.layers)):
            seq_enc, _ = self.layers[i](seq_enc, mask)

        return enc_seq, seq_enc
    
    
class Decoder(nn.Module):
    
    def __init__(self, decoder_type, d_model, all_size, n_head, d_inner_hid, d_k, d_v, mask, dropout, device):
        super().__init__()

        self.all_size = all_size
        self.num_heads = n_head
        self.mask = mask
        self.decoder_type = decoder_type

        self.indexes = last_points(self.all_size, device)
        if decoder_type == "Attention":
            self.layers = nn.ModuleList([
                DecoderLayer(d_model, d_inner_hid, n_head, d_k, d_v, dropout, False),
                DecoderLayer(d_model, d_inner_hid, n_head, d_k, d_v, dropout, False)
            ])
    
    
    def forward(self, seq_enc, refer):
        if self.decoder_type == "Attention":
            mask = self.mask.repeat(len(seq_enc), self.num_heads, 1, 1).to(seq_enc.device)
            dec_enc, _  = self.layers[0](seq_enc, refer, refer, mask)
            #refer_enc = torch.cat([refer, dec_enc], dim=1)
            dec_enc, _ = self.layers[1](dec_enc, refer, refer, mask)
        else:
            dec_enc = refer
        indexes = self.indexes.repeat(dec_enc.size(0), 1, dec_enc.size(2)).to(dec_enc.device)
        all_enc = torch.gather(dec_enc, 1, indexes)
        dec_enc = all_enc.view(dec_enc.size(0), -1)
        # indexes = indexes.view(dec_enc.size(0), -1, dec_enc.size(2))
        # all_enc = torch.gather(dec_enc, 1, indexes)
        # dec_enc = all_enc.view(dec_enc.size(0), self.all_size[0], -1)
        return dec_enc

class Predictor(nn.Module):

    def __init__(self, dim, output_size=1):
        super().__init__()

        self.linear= nn.Linear(dim,  output_size, bias=False)
        self.softplus = nn.Sigmoid()

        nn.init.xavier_normal_(self.linear.weight)

    def forward(self, data):
        out = self.linear(data)
        out = self.softplus(out)
        return out


class Pyraformer(nn.Module):

    def __init__(self, d_model, window_set, input_size, feature_size, inner_size, n_layer, n_head, 
                 d_inner_hid, d_k, d_v, num_tokens, decoder_type, dropout, device):
        super().__init__()
        assert decoder_type in ["Attention", "FC"]
        mask, all_size = get_mask(input_size, window_set, inner_size, device)

        self.encoder = Encoder(d_model, window_set, input_size, feature_size, n_layer, n_head, d_inner_hid, 
                               d_k, d_v, num_tokens, mask, dropout, device)

        self.decoder = Decoder(decoder_type, d_model, all_size, n_head, d_inner_hid, d_k, d_v, mask, dropout, device)
        self.preditor = Predictor(d_model * len(all_size), 1)


        # convert hidden vectors into two scalar

    def forward(self, data, token=None):
        enc_seq, enc_out = self.encoder(data, token)
        dec_output = self.decoder(enc_seq, enc_out)

        out = self.preditor(dec_output)

        return out

